@sfmig your comment on #5 indicates that it would be useful to have certain bits of information from the plot as outputs from this kind of function. Currently I'm just exposing the other hist2d outputs (that are suppressed by the wrapper otherwise) to the user here, not sure if you had more detailed outputs in mind when writing your comment.

But if so, we can also close #5 with this PR too.

thanks for checking @willGraham01 !

The point of that comment was that often you not only want the figure, but also the 2D array with the bin counts. From your comment ...

Currently I'm just exposing the other hist2d outputs

seems like that is covered? So I think we can close #5 yay 😄 🚀

(Just fyi I vaguely remember this was something Sepi requested but not sure)

(Just fyi I vaguely remember this was something Sepi requested but not sure)

I hope it is b/c otherwise I've just wasted 5 hours of Niko's grant 🤭 😂 But will mark #5 as closable by this 🥳

I will finish reviewing this tomorrow, but I can already do some cool things with this!

import numpy as np
from matplotlib import pyplot as plt

from movement import sample_data
from movement.plots import plot_occupancy

# Load the sample dataset 
ds = sample_data.fetch_dataset("DLC_two-mice.predictions.csv")

# Compute the centroid of all keypoints
centroid_position = ds.position.mean("keypoints")

image = plt.imread(ds.attrs["frame_path"])

# Construct bins of size 20x20 pixels that cover the entire image
bin_pix = 30
bins = [
    np.arange(0, image.shape[0] + bin_pix, bin_pix),
    np.arange(0, image.shape[1] + bin_pix, bin_pix),
]

# Initialize the figure and axis
fig, ax = plt.subplots()

# Show the image
ax.imshow(image)

# Plot the occupancy 2D histogram for each individual
_, _, hist_data = plot_occupancy(
    da=centroid_position,
    selection={"individuals": "individual1"},
    ax=ax,
    cmap="viridis",
    alpha=0.5,
    bins=bins,
    cmin=3,      # Set the minimum shown count
    norm="log"
)

# Set the axis limits to match the image
ax.set_xlim(0, image.shape[1])
ax.set_ylim(image.shape[0], 0)

Function signature is now in line with plot trajectory, and I've implemented the default behaviour of "take centroid of keypoints, then aggregate over individuals".

Note that the aggregation over individuals works by "stacking" the input DataArrays "individual" axis along the "time" axis. EG a (10, 2, 4) time-space-individuals array can be reshaped and counted as a (40, 2) array. This:

• Makes it robust against NaN values (NaN values can be filtered on a per-position basis, since xarray.dropna currently doesn't support dropping NaNs along multiple axes simultaneously).
• Prevents any issues with bin sizes not aligning. If we count each individual first, then sum the counts, providing bins=[30,20] will give non-aligned bin edges. By stacking first, we avoid this issue (note that passing explicit bin edges works in both cases, regardless).